#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
测试向量恢复功能

本测试文件专门测试知识处理器中的向量恢复功能，
确保系统重启后能够正确从ChromaDB恢复向量数据到内存向量存储。
"""

import asyncio
import pytest
import pytest_asyncio
import sys
import os
from pathlib import Path
from loguru import logger

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from src.core.knowledge_processor import KnowledgeProcessor
from src.database.manager import DatabaseManager
from src.config import load_config


class Test向量恢复功能:
    """测试向量恢复功能的测试类"""
    
    @pytest_asyncio.fixture
    async def setup_knowledge_processor(self):
        """设置知识处理器"""
        try:
            # 加载配置
            config = load_config()
            
            # 创建数据库管理器
            db_manager = DatabaseManager(config)
            await db_manager.initialize()
            
            # 创建知识处理器
            knowledge_processor = KnowledgeProcessor(db_manager, config)
            await knowledge_processor.initialize()
            
            yield knowledge_processor
            
            # 清理
            await knowledge_processor.shutdown()
            await db_manager.shutdown()
            
        except Exception as e:
            logger.error(f"设置知识处理器失败: {e}")
            raise
    
    @pytest.mark.asyncio
    async def test_向量存储初始状态(self, setup_knowledge_processor):
        """测试向量存储的初始状态"""
        knowledge_processor = setup_knowledge_processor
        
        # 检查vectorizer是否正确初始化
        assert knowledge_processor.vectorizer is not None, "vectorizer未正确初始化"
        assert knowledge_processor.vectorizer.vector_store is not None, "vector_store未正确初始化"
        
        # 检查向量存储类型
        logger.info(f"向量存储类型: {type(knowledge_processor.vectorizer.vector_store)}")
        
        # 检查初始向量数量
        if hasattr(knowledge_processor.vectorizer.vector_store, 'vectors'):
            initial_count = len(knowledge_processor.vectorizer.vector_store.vectors)
            logger.info(f"初始向量数量: {initial_count}")
        
    @pytest.mark.asyncio
    async def test_chromadb连接状态(self, setup_knowledge_processor):
        """测试ChromaDB连接状态"""
        knowledge_processor = setup_knowledge_processor
        
        # 检查ChromaDB客户端
        assert knowledge_processor.chromadb_client is not None, "ChromaDB客户端未初始化"
        
        # 检查documents集合
        try:
            documents_collection = knowledge_processor.chromadb_client.get_collection('documents')
            assert documents_collection is not None, "documents集合不存在"
            
            # 获取集合中的文档数量
            count = documents_collection.count()
            logger.info(f"ChromaDB中documents集合文档数量: {count}")
            
        except Exception as e:
            logger.error(f"访问ChromaDB集合失败: {e}")
            raise
    
    @pytest.mark.asyncio
    async def test_向量恢复逻辑触发(self, setup_knowledge_processor):
        """测试向量恢复逻辑是否正确触发"""
        knowledge_processor = setup_knowledge_processor
        
        # 清空内存向量存储（模拟系统重启后的状态）
        if hasattr(knowledge_processor.vectorizer.vector_store, 'vectors'):
            original_vectors = knowledge_processor.vectorizer.vector_store.vectors.copy()
            knowledge_processor.vectorizer.vector_store.vectors.clear()
            knowledge_processor.vectorizer.vector_store.metadatas.clear()
            
            logger.info("已清空内存向量存储，模拟系统重启状态")
            
            # 执行搜索，应该触发向量恢复
            try:
                results = await knowledge_processor.search_knowledge(
                    query="Python编程",
                    user_id="test_user",
                    limit=5
                )
                
                logger.info(f"搜索结果数量: {len(results)}")
                
                # 检查向量是否已恢复
                restored_count = len(knowledge_processor.vectorizer.vector_store.vectors)
                logger.info(f"恢复后向量数量: {restored_count}")
                
                # 恢复原始向量（清理）
                knowledge_processor.vectorizer.vector_store.vectors = original_vectors
                
            except Exception as e:
                logger.error(f"向量恢复测试失败: {e}")
                # 恢复原始向量（清理）
                knowledge_processor.vectorizer.vector_store.vectors = original_vectors
                raise
    
    @pytest.mark.asyncio
    async def test_向量恢复方法直接调用(self, setup_knowledge_processor):
        """直接测试向量恢复方法"""
        knowledge_processor = setup_knowledge_processor
        
        # 记录原始向量数量
        if hasattr(knowledge_processor.vectorizer.vector_store, 'vectors'):
            original_count = len(knowledge_processor.vectorizer.vector_store.vectors)
            logger.info(f"原始向量数量: {original_count}")
            
            # 清空向量存储
            knowledge_processor.vectorizer.vector_store.vectors.clear()
            knowledge_processor.vectorizer.vector_store.metadatas.clear()
            
            # 直接调用恢复方法
            try:
                await knowledge_processor._restore_vectors_from_chromadb()
                
                # 检查恢复结果
                restored_count = len(knowledge_processor.vectorizer.vector_store.vectors)
                logger.info(f"恢复后向量数量: {restored_count}")
                
                assert restored_count > 0, "向量恢复失败，恢复数量为0"
                
            except Exception as e:
                logger.error(f"直接调用向量恢复方法失败: {e}")
                raise
    
    @pytest.mark.asyncio
    async def test_搜索功能完整性(self, setup_knowledge_processor):
        """测试搜索功能的完整性"""
        knowledge_processor = setup_knowledge_processor
        
        # 执行多个不同的搜索查询
        test_queries = [
            "Python编程",
            "数据库",
            "机器学习",
            "web开发",
            "算法"
        ]
        
        for query in test_queries:
            try:
                results = await knowledge_processor.search_knowledge(
                    query=query,
                    user_id="test_user",
                    limit=3
                )
                
                logger.info(f"查询'{query}'返回{len(results)}个结果")
                
                # 检查结果格式
                for result in results:
                    assert 'content' in result, f"搜索结果缺少content字段: {result}"
                    assert 'score' in result, f"搜索结果缺少score字段: {result}"
                    
            except Exception as e:
                logger.error(f"搜索查询'{query}'失败: {e}")
                # 不抛出异常，继续测试其他查询


if __name__ == "__main__":
    # 配置日志
    logger.remove()
    logger.add(
        sys.stdout,
        format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>",
        level="INFO"
    )
    
    # 运行测试
    pytest.main([__file__, "-v", "-s"])